import time

import matplotlib.pyplot as plt
import numpy as np

from cion.emccd.camera import AcquisitionMode, EmccdCameraByDll, ReadMode, ShutterMode, ShutterTTLSignal, Status, \
    create_uint32_buffer, EmccdDll
from cion.emccd.functions import gaussian_2d, gaussian, calibrate_emccd_fit, emccd_fit_once, find_ions,counts_fit
from os import getcwd, makedirs
from os.path import exists
from scipy.optimize import curve_fit


class EMCCDForIon:
    def __init__(self):
        self.dll_file_path = "cion/emccd/atmcd64d.dll"
        self.target_temperature = -80
        self.camera = EmccdCameraByDll(EmccdDll(self.dll_file_path))

        # temperature
        self.camera.set_temperature(-80)
        self.camera.set_cooler(True)

        # shutter
        self.camera.set_shutter(ShutterTTLSignal.LOW, ShutterMode.PERMANENTLY_OPEN, 27, 27)

        # baseline clamp
        if self.camera.support_baseline_clamp:
            self.camera.set_baseline_clamp(True)
        if self.camera.support_baseline_offset:
            self.camera.set_baseline_offset(0)

        # set shift speeds & ad channels
        self.camera.set_hs_speed(0, 0)
        self.camera.set_vs_speed(4)
        # self.camera.dll.SetVSAmplitude(0)
        # self.camera.dll.SetFrameTransferMode(0)

        # pre-amp gain
        self.camera.set_pre_amp_gain(2)

        # em & em gain
        self.camera.set_em_advanced(True)
        self.camera.set_em_gain(1000)

        # set read mode & set image
        self.camera.set_read_mode(ReadMode.IMAGE)
        self.img_w, self.img_h = self.camera.get_detector_size()
        self.camera.set_image(hbin=1, vbin=1, hstart=0, hend=self.img_w, vstart=0, vend=self.img_h)
        self.img_size = self.img_w * self.img_h

        #cooling
        while True:
            temperature = self.camera.get_temperature()
            if temperature > self.target_temperature:
                print(f"Still cooling ...  (t = {temperature} ℃)")
                time.sleep(1)
            else:
                print(f"Reached target temperature ...  (t = {temperature} ℃)")
                break
        self.img_num = 1
        # set acquisition mode
        self.set_acquisition_mode()
        self.calibrate_parameter=None
        # TODO fast kinetic 可能还有别的参数要设置

    def set_acquisition_mode(self, mode=AcquisitionMode.KINETICS):
        self.camera.set_acquisition_mode(mode)
        self.camera.dll.SetNumberKinetics(self.img_num)
        self.camera.set_exposure_time(0.001)
        self.camera.dll.SetAccumulationCycleTime(0.1)
        self.camera.set_kinetic_cycle_time(0.1)
        exposure_time, accumulate_cycle_time, kinetic_cycle_time = self.camera.get_acquisition_times()
        self.camera.dll.SetTriggerMode(0)

    def set_exposure_time(self, exposuretime):
        self.camera.set_exposure_time(exposuretime)

    def set_image(self, hbin=1, vbin=1, hstart=0, vstart=0):
        self.camera.set_image(hbin=hbin, vbin=vbin, hstart=hstart, hend=self.img_w, vstart=vstart, vend=self.img_h)

    def set_img_num(self, img_num):
        self.img_num = img_num
        self.camera.dll.SetNumberKinetics(self.img_num)

    def measure_and_fit(self):
        self.set_img_num(1)
        buffer = create_uint32_buffer(self.img_size * self.img_num)

        # capture
        self.camera.start_acquisition()
        while True:
            time.sleep(1)
            status = self.camera.get_status()
            if status == Status.IDLE:
                break
        self.camera.get_acquired_data(self.img_size * self.img_num, buffer)
        # parse data as image(s)
        images = np.frombuffer(buffer, dtype=np.uint32).reshape((-1, self.img_h, self.img_w))

        # 获取此py文件路径，在此路径选创建在new_folder文件夹中的test文件夹
        folder = getcwd() + '\\measure_image\\'
        if not exists(folder):
            makedirs(folder)
        # show image(s)
        i = 0
        plt.clf()
        for image in images:
            plt.figure()
            plt.imshow(image)
            i += 1
            plt.savefig(folder + 'measure%d' % i + '.png')
            plt.clf()
        ion_position_list=find_ions(np.uint32(images[0]))
        ion_position = [ion_position_list[0,0],ion_position_list[0,1]]
        ion_size = ion_position_list[0,2]
        ion_image = images[0][ion_position[0] - ion_size:ion_position[0] + ion_size,
                    ion_position[1] - ion_size:ion_position[1] + ion_size]
        ion_image = ion_image - np.median(ion_image)

        plt.imshow(ion_image)
        plt.savefig(folder + 'find_ion_position.png')
        plt.clf()
        count_matrix_1d = np.sum(ion_image, axis=0)
        plt.plot(count_matrix_1d)
        plt.savefig(folder + 'count_matrix_1d.png')
        plt.clf()
        x = range(len(ion_image))
        return emccd_fit_once(x, ion_image, {'mu': self.calibrate_parameter['mu'], 'sigma': self.calibrate_parameter['sigma']})
        # print("emccd_fit的结果是",)

    def calibrate(self):
        # set acquisition mode
        self.set_img_num(100)
        # Take data
        # prepare buffer
        buffer = create_uint32_buffer(self.img_size * self.img_num)
        # capture
        self.camera.start_acquisition()
        while True:
            time.sleep(1)
            status = self.camera.get_status()
            if status == Status.IDLE:
                break
        self.camera.get_acquired_data(self.img_size * self.img_num, buffer)
        # parse data as image(s)
        images = np.frombuffer(buffer, dtype=np.uint32).reshape((-1, self.img_h, self.img_w))
        image_averaged = sum(images) / len(images)

        folder = getcwd() + '\\calibrate_image\\'
        if not exists(folder):
            makedirs(folder)
        plt.clf()
        plt.figure()
        plt.imshow(image_averaged)
        plt.savefig(folder + 'image_averaged' + '.png')
        plt.clf()
        ion_position_list = find_ions(np.uint32(image_averaged))
        ion_position = [ion_position_list[0, 0], ion_position_list[0, 1]]
        ion_size = ion_position_list[0,2]
        ion_image = image_averaged[ion_position[0] - ion_size:ion_position[0] + ion_size,
                    ion_position[1] - ion_size:ion_position[1] + ion_size]
        ion_image = ion_image - np.median(ion_image)
        plt.imshow(ion_image)
        plt.savefig(folder + 'ion_image' + '.png')
        plt.clf()
        count_matrix_1d = np.sum(ion_image, axis=0)

        plt.plot(count_matrix_1d)
        plt.savefig(folder + 'calibrate_count_matrix_1d' + '.png')
        plt.clf()
        def calibrate_EMCCD_fit(count_matrix,path):
            y = np.sum(count_matrix, axis=0)
            x = range(len(y))
            bounds = (np.array([0, 0, 0, -np.inf]), [np.inf, np.inf, np.inf, np.inf])
            params_cali, cov = curve_fit(gaussian, x, y, p0=[np.median(x), len(x) / 2, 100, 0], bounds=bounds)

            x_fit = np.arange(x[0], x[-1], (x[-1] - x[0]) / 50)
            plt.figure()
            plt.plot(x, count_matrix_1d, 'o')
            plt.plot(x_fit, gaussian(x_fit, params_cali[0], params_cali[1], params_cali[2], params_cali[3]))
            plt.savefig(folder + 'calibrate_emccd_fit' + '.png')
            plt.clf()
            return {'mu': params_cali[0], 'sigma': params_cali[1], 'amp': params_cali[2], 'bias': params_cali[3]}

        params_cali = calibrate_EMCCD_fit(ion_image,path=folder)
        self.calibrate_parameter = params_cali
        amp_list = []
        i=0
        for image in images:
            images_single_shot = image[ion_position[0] - ion_size:ion_position[0] + ion_size,
                                 ion_position[1] - ion_size:ion_position[1] + ion_size]
            images_single_shot = images_single_shot - np.median(images_single_shot)
            x = range(len(images_single_shot))
            amp_list = amp_list + [emccd_fit_once(x, images_single_shot, params_cali)]
            count_matrix_1d = np.sum(images_single_shot, axis=0)
            plt.plot(count_matrix_1d)
            plt.savefig(folder + 'calibrate_count_matrix_1d%d' % i + '.png')
            i=i+1
            plt.clf()
        plt.hist(amp_list)
        plt.savefig(folder + 'calibrate_amp_list.png')
        plt.clf()

        ion_position_list = find_ions(np.uint32(image_averaged))
        ion_position = [ion_position_list[0, 0], ion_position_list[0, 1]]
        ion_size = ion_position_list[0, 2]
        ion_image_dark = image_averaged[ion_position[0] - ion_size:ion_position[0] + ion_size,
                         ion_position[1] - ion_size:ion_position[1] + ion_size]
        ion_image_dark = ion_image_dark - np.median(ion_image_dark)

        plt.imshow(ion_image_dark)
        plt.savefig(folder + 'calibrate_ion_image_dark.png')
        plt.clf()

        amp_list_dark = []
        for image in images:
            images_single_shot = image[ion_position[0] - ion_size:ion_position[0] + ion_size,
                                 ion_position[1] - ion_size:ion_position[1] + ion_size]
            images_single_shot = images_single_shot - np.median(images_single_shot)
            x = range(len(images_single_shot))
            amp_list_dark = amp_list_dark + [emccd_fit_once(x, images_single_shot, params_cali)]


        plt.hist(amp_list_dark)
        plt.savefig(folder + 'calibrate_amp_list_dark.png')
        plt.clf()
        plt.hist(amp_list_dark, 20)
        plt.hist(amp_list, 40)
        plt.savefig(folder + 'calibrate_amp_list_dark_two_combined.png')
        plt.clf()
        threshold_list, infidelity_list, threshold, infidelity = counts_fit.get_infidelity(amp_list_dark, amp_list,
                                                                                range(0, 1000, 10))
        plt.plot(threshold_list, infidelity_list)
        plt.savefig(folder + 'threshold_infidelity.png')
        plt.clf()
